from numpy import (
    ndarray,  # ndarray 类型
    array,  # 创建数组
    reshape,  # 改变数组形状
    arange,  # 创建范围数组
    zeros,  # 创建全零数组
    ones,  # 创建全一数组
    arccos,  # 计算反余弦
    pi,  # 圆周率
)

from src.config.config import (
    Config,  # 配置类
)
from src.data.satellite import (
    Satellite,  # 卫星类
)


class SatelliteManager:
    """
    SatelliteManager 负责管理所有的卫星对象及其辅助功能。
    """

    def __init__(
            self,
            config: Config,
    ) -> None:
        """
        初始化 SatelliteManager 类。

        参数:
        - config: 配置对象，包含生成随机数和日志记录器的配置。
        """

        self.rng = config.rng  # 随机数生成器
        self.logger = config.logger.getChild(__name__)  # 日志记录器

        self.satellites: list[Satellite] = []  # 卫星对象列表
        self._initialize_satellites(config=config)  # 初始化卫星对象

        self.channel_state_information: ndarray = array([])  # 信道状态信息矩阵
        self.erroneous_channel_state_information: ndarray = array([])  # 带有误差的信道状态信息矩阵

        self.logger.info('卫星设置完成')

    def calc_spherical_coordinates(
            self,
            config,
    ) -> ndarray:
        """
        计算卫星的球面坐标，包括半径、倾斜角和离开角（AOD）。

        参数:
        - config: 配置对象，包含卫星数量、距离等信息。

        返回:
        - ndarray: 卫星的球面坐标数组，形状为 (3, 卫星数量)。
        """

        # 计算平均卫星位置
        sat_pos_average = (arange(0, config.sat_nr) - (config.sat_nr - 1) / 2) * config.sat_dist_average

        # 在卫星距离上添加随机值
        random_factor = self.rng.uniform(low=-config.sat_dist_bound,
                                         high=config.sat_dist_bound,
                                         size=config.sat_nr)
        sat_dist = sat_pos_average + random_factor

        # 计算卫星的AOD偏差
        sat_aods_diff_earth_rad = zeros(config.sat_nr)

        for sat_idx in range(config.sat_nr):
            if sat_dist[sat_idx] < 0:
                sat_aods_diff_earth_rad[sat_idx] = -1 * arccos(1 - 0.5 * (sat_dist[sat_idx] / config.radius_orbit)**2)
            elif sat_dist[sat_idx] >= 0:
                sat_aods_diff_earth_rad[sat_idx] = arccos(1 - 0.5 * (sat_dist[sat_idx] / config.radius_orbit)**2)

        # 计算卫星中心的AOD
        sat_center_aod_earth_rad = config.sat_center_aod_earth_deg * pi / 180

        # 计算卫星的AOD
        sat_aods_earth_rad = sat_center_aod_earth_rad + sat_aods_diff_earth_rad

        # 创建卫星的球面坐标
        sat_radii = config.radius_orbit * ones(config.sat_nr)  # 卫星轨道半径
        sat_inclinations = pi / 2 * ones(config.sat_nr)  # 卫星倾斜角

        sat_spherical_coordinates = array([sat_radii, sat_inclinations, sat_aods_earth_rad])

        return sat_spherical_coordinates

    def _initialize_satellites(
            self,
            config: Config,
    ) -> None:
        """
        根据配置初始化卫星对象列表。

        参数:
        - config: 配置对象，包含卫星的相关参数。
        """

        sat_spherical_coordinates = self.calc_spherical_coordinates(config=config)

        for sat_idx in range(config.sat_nr):
            self.satellites.append(
                Satellite(
                    idx=sat_idx,
                    spherical_coordinates=sat_spherical_coordinates[:, sat_idx],
                    **config.satellite_args,
                )
            )

    def update_positions(
            self,
            config,
    ) -> None:
        """
        更新卫星的位置。

        参数:
        - config: 配置对象，包含更新卫星位置所需的参数。
        """

        sat_spherical_coordinates = self.calc_spherical_coordinates(config=config)

        for satellite in self.satellites:
            satellite.update_position(
                spherical_coordinates=sat_spherical_coordinates[:, satellite.idx],
            )

    def calculate_satellite_distances_to_users(
            self,
            users: list,
    ) -> None:
        """
        计算每个卫星与每个用户之间的距离。

        参数:
        - users: 用户对象列表。
        """

        for satellite in self.satellites:
            satellite.calculate_distance_to_users(users=users)

    def calculate_satellite_aods_to_users(
            self,
            users: list,
    ) -> None:
        """
        计算每个卫星到每个用户的离开角（AOD）。

        参数:
        - users: 用户对象列表。
        """

        for satellite in self.satellites:
            satellite.calculate_aods_to_users(users=users)

    def calculate_steering_vectors_to_users(
            self,
            users: list,
    ) -> None:
        """
        计算每个卫星到每个用户的引导矢量（每个天线一个值）。

        参数:
        - users: 用户对象列表。
        """

        for satellite in self.satellites:
            satellite.calculate_steering_vectors(users=users)

    def update_channel_state_information(
            self,
            channel_model,
            users: list,
    ) -> None:
        """
        更新每个卫星天线与用户之间的信道状态信息，并将其整合到全局信道状态信息矩阵中。

        参数:
        - channel_model: 信道模型，用于计算信道状态信息。
        - users: 用户对象列表。
        """

        # 更新每个卫星的信道状态信息
        for satellite in self.satellites:
            satellite.update_channel_state_information(channel_model=channel_model, users=users)

        # 构建全局信道状态信息
        channel_state_per_satellite = zeros((len(users), self.satellites[0].antenna_nr, len(self.satellites)),
                                            dtype='complex')
        for satellite in self.satellites:
            channel_state_per_satellite[:, :, satellite.idx] = satellite.channel_state_to_users
        self.channel_state_information = reshape(
            channel_state_per_satellite, (len(users), self.satellites[0].antenna_nr * len(self.satellites)))

    def update_erroneous_channel_state_information(
            self,
            error_model_config,
            users: list,
    ) -> None:
        """
        更新每个卫星的带有误差的信道状态信息，并将其整合到全局信道状态信息矩阵中。

        参数:
        - error_model_config: 误差模型配置，用于生成带有误差的信道状态信息。
        - users: 用户对象列表。
        """

        # 应用误差模型更新每个卫星的信道状态信息
        for satellite in self.satellites:
            satellite.update_erroneous_channel_state_information(error_model_config=error_model_config, users=users)

        # 收集全局带有误差的信道状态信息
        erroneous_channel_state_per_satellite = zeros(
            (len(users), self.satellites[0].antenna_nr, len(self.satellites)),
            dtype='complex',
        )
        for satellite in self.satellites:
            erroneous_channel_state_per_satellite[:, :, satellite.idx] = satellite.erroneous_channel_state_to_users
        self.erroneous_channel_state_information = reshape(
            erroneous_channel_state_per_satellite, (len(users), self.satellites[0].antenna_nr * len(self.satellites)))

    def get_aods_to_users(
            self,
    ) -> ndarray:
        """
        获取每个卫星到每个用户的离开角（AOD）。

        返回:
        - ndarray: 每个卫星到每个用户的离开角数组。
        """

        aods_to_users = zeros((len(self.satellites), len(self.satellites[0].aods_to_users)))
        for satellite_id, satellite in enumerate(self.satellites):
            aods_to_users[satellite_id, :] = array(list(satellite.aods_to_users.values()))

        return aods_to_users
